package com.joker;

import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;
import java.io.IOException;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.CountDownLatch;

public class DistributedLockNative {
    private ZooKeeper zk;
    private String lockPath;
    private String currentLock;
    private String waitPath;
    private CountDownLatch connectedLatch = new CountDownLatch(1);
    private CountDownLatch waitLatch = new CountDownLatch(1);
    
    private static final String ZK_ADDRESS = "localhost:2181";
    private static final int SESSION_TIMEOUT = 5000;
    private static final String LOCK_ROOT = "/locks";
    private static final String LOCK_PREFIX = LOCK_ROOT + "/lock-";

    public DistributedLockNative() throws IOException, InterruptedException {
        this.zk = new ZooKeeper(ZK_ADDRESS, SESSION_TIMEOUT, new Watcher() {
            @Override
            public void process(WatchedEvent event) {
                if (event.getState() == Event.KeeperState.SyncConnected) {
                    connectedLatch.countDown();
                }
                
                // 监听前一个节点的删除事件
                if (event.getType() == Event.EventType.NodeDeleted && 
                    event.getPath().equals(waitPath)) {
                    waitLatch.countDown();
                }
            }
        });
        connectedLatch.await();
        
        // 确保根节点存在
        try {
            Stat stat = zk.exists(LOCK_ROOT, false);
            if (stat == null) {
                zk.create(LOCK_ROOT, new byte[0], ZooDefs.Ids.OPEN_ACL_UNSAFE, 
                         CreateMode.PERSISTENT);
            }
        } catch (KeeperException e) {
            // 节点已存在，忽略
        }
    }

    // 获取锁
    public boolean lock() {
        try {
            // 创建临时顺序节点
            currentLock = zk.create(LOCK_PREFIX, new byte[0], 
                    ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.EPHEMERAL_SEQUENTIAL);
            
            // 获取所有锁节点
            List<String> lockNodes = zk.getChildren(LOCK_ROOT, false);
            Collections.sort(lockNodes);
            
            // 当前节点是最小的节点，获取锁成功
            if (currentLock.equals(LOCK_ROOT + "/" + lockNodes.get(0))) {
                System.out.println(Thread.currentThread().getName() + " 获取锁成功: " + currentLock);
                return true;
            }
            
            // 不是最小节点，找到前一个节点并监听
            int currentIndex = lockNodes.indexOf(currentLock.substring(LOCK_ROOT.length() + 1));
            String prevNode = lockNodes.get(currentIndex - 1);
            waitPath = LOCK_ROOT + "/" + prevNode;
            
            // 检查前一个节点是否存在，并设置监听
            Stat stat = zk.exists(waitPath, true);
            if (stat != null) {
                System.out.println(Thread.currentThread().getName() + " 等待锁: " + waitPath);
                waitLatch.await(); // 等待前一个节点被删除
            }
            
            System.out.println(Thread.currentThread().getName() + " 获取锁成功: " + currentLock);
            return true;
            
        } catch (Exception e) {
            e.printStackTrace();
            return false;
        }
    }

    // 释放锁
    public void unlock() {
        try {
            if (currentLock != null) {
                zk.delete(currentLock, -1);
                System.out.println(Thread.currentThread().getName() + " 释放锁: " + currentLock);
                currentLock = null;
            }
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public void close() throws InterruptedException {
        if (zk != null) {
            zk.close();
        }
    }
}